from args import *
from elasticsearch import Elasticsearch, helpers
import math
client = Elasticsearch(
            hosts=ELASTIC_ADDRESS, basic_auth=(ELASTIC_USERNAME, ELASTIC_PASSWORD),
            timeout=10000)

def get_query(from_, size, start_date=None, end_date=None):
    query = {
        "query": {
            "range": {
                "insert_time": {
                            "gte": start_date,
                        "lte": end_date
                        }
            }
        },
        "from": from_,
        "size": size
    }
    return query
def get_count(start_date: str, end_date: str):
    query = wrap_read_query(start_date, end_date)
    return client.count(index=INDEX_NAME, body=query)['count']

def wrap_read_query(start_date: str, end_date: str):
    query = {
        "query": {
            "bool": {
                "filter": {
                    "range": {
                        "insert_time": {
                            "lte": end_date,
                            "gte": start_date
                        }
                    }
                }
            }
        }
    }
    return query

def batch_read(start_date: str, end_date: str):
    start_date = start_date + " 00:00:00"
    end_date = end_date + " 23:59:59"
    count = get_count(start_date, end_date)
    batch_size = BATCH_SIZE_ES
    epochs = int(math.ceil(count / batch_size))
    ret = []
    for i in range(epochs):
        read_query = wrap_read_query(start_date, end_date)
        from_ = i * batch_size
        size = min(batch_size, count - from_)
        read_query["from"] = from_
        read_query["size"] = size
        response = client.search(index=INDEX_NAME, body=read_query)
        body = response.body
        hits = body['hits']['hits']
        ret += ([hit['_source'] for hit in hits])
    return ret